import numpy as np
import matplotlib.pyplot as plt

# ---------------------- 参数设置 ----------------------
t0 = 0.0       # 激励起始时间 (s)
w = 6.28  # 基频角速度 (rad/s)，若频率f=1Hz，则w=2πf
A0 = 1.0       # 静态偏移量
N = 1         # 谐波次数（n=1,2,...,N）
An = [0.8]  # 各次谐波的余弦项振幅（A₁, A₂, ..., Aₙ）
Bn = [1.9]  # 各次谐波的正弦项振幅（B₁, B₂, ..., Bₙ）
t_start = 0.0  # 时间范围起始
t_end = 1.0    # 时间范围结束
dt = 0.05      # 时间步长

# 生成时间数组
t = np.linspace(t_start, t_end, int((t_end - t_start) / dt) + 1)
tau = t - t0  # 时间偏移量（t - t0）

# ---------------------- 计算位移u(t) ----------------------
u = np.full_like(t, A0)  # 初始化为静态偏移A0
mask = t < t0            # 标记激励前的区域（t < t0）
for n in range(1, N+1):
    n_idx = n - 1  # 数组索引从0开始
    cos_term = An[n_idx] * np.cos(n * w * tau)
    sin_term = Bn[n_idx] * np.sin(n * w * tau)
    # 仅在t >= t0时叠加谐波项
    u[~mask] += cos_term[~mask] + sin_term[~mask]

# ---------------------- 计算速度v(t)（u的一阶导数） ----------------------
v = np.zeros_like(t)
for n in range(1, N+1):
    n_idx = n - 1
    # 余弦项导数：d/dt [Aₙcos(nωτ)] = -Aₙnω sin(nωτ)
    # 正弦项导数：d/dt [Bₙsin(nωτ)] = Bₙnω cos(nωτ)
    sin_deriv = -An[n_idx] * n * w * np.sin(n * w * tau)
    cos_deriv = Bn[n_idx] * n * w * np.cos(n * w * tau)
    v[~mask] += sin_deriv[~mask] + cos_deriv[~mask]

# ---------------------- 计算加速度a(t)（u的二阶导数） ----------------------
a = np.zeros_like(t)
for n in range(1, N+1):
    n_idx = n - 1
    # 余弦项二阶导数：d²/dt² [Aₙcos(nωτ)] = -Aₙn²ω² cos(nωτ)
    # 正弦项二阶导数：d²/dt² [Bₙsin(nωτ)] = -Bₙn²ω² sin(nωτ)
    cos_2deriv = -An[n_idx] * (n**2) * (w**2) * np.cos(n * w * tau)
    sin_2deriv = -Bn[n_idx] * (n**2) * (w**2) * np.sin(n * w * tau)
    a[~mask] += cos_2deriv[~mask] + sin_2deriv[~mask]

# ---------------------- 可视化结果 ----------------------
plt.figure(figsize=(10, 8))

# 位移曲线
plt.subplot(3, 1, 1)
plt.plot(t, u, label='$A(t)$')
plt.axvline(x=t0, color='r', linestyle='--', label=f'$t_0={t0}$')
plt.xlabel('time $t$ (s)')
plt.ylabel('displacement')
plt.legend()
plt.grid(True)

# 速度曲线
plt.subplot(3, 1, 2)
plt.plot(t, v, label='$dA/dt(t)$', color='g')
plt.axvline(x=t0, color='r', linestyle='--')
plt.xlabel('time $t$ (s)')
plt.ylabel('velocity')
plt.legend()
plt.grid(True)

# 加速度曲线
plt.subplot(3, 1, 3)
plt.plot(t, a, label='$d²A/dt²(t)$', color='orange')
plt.axvline(x=t0, color='r', linestyle='--')
plt.xlabel('time $t$ (s)')
plt.ylabel('acceleration')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()